# -*- coding: utf-8 -*- 
# @Time : 2022/4/2 15:51 
# @Author : zzuxyj 
# @File : 07-nn-maxPool.py

"""
最大池化操作

会计算池化之后的大小
"""
import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d
from torch.utils.data import DataLoader

# 加载数据
from torch.utils.tensorboard import SummaryWriter

test_set = torchvision.datasets.CIFAR10("../dataset/CIFAR10", download=True,
                                        transform=torchvision.transforms.ToTensor())
dataloader = DataLoader(dataset=test_set, batch_size=64)

# tensorboard
writer = SummaryWriter("logs07")

# 定义网络模型
class Model(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.conv01 = Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0) # (32-3+1)/1=30
        self.maxPool = MaxPool2d(kernel_size=3, ceil_mode=True) # (30-3+3)/3 = 10  channel=6

    def forward(self, input):
        output = self.conv01(input)
        output = self.maxPool(output)
        return output


# 声明模型对象
model = Model()

# 遍历数据集
step = 0
for data in dataloader:
    imgs , target = data
    imgs = model(imgs)
    imgs = torch.reshape(imgs , (-1,3,10,10))  # batch_size=64 -》 batch_size=128
    writer.add_images("maxPool" , imgs , global_step=step)
    step+=1

# 关闭
writer.close()
